"""
主机信息自动获取模块
通过SSH连接获取主机的系统信息
"""
import paramiko
from io import StringIO


def get_host_info(ip_address, port, username, password, private_key=None):
    """
    通过SSH获取主机信息
    """
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    try:
        # 连接SSH
        if private_key:
            key = paramiko.RSAKey.from_private_key(StringIO(private_key))
            ssh.connect(hostname=ip_address, port=port, username=username, pkey=key, timeout=10)
        else:
            ssh.connect(hostname=ip_address, port=port, username=username, password=password, timeout=10)
        
        # 获取操作系统信息
        os_type = get_os_info(ssh)
        
        # 获取CPU核心数
        cpu_cores = get_cpu_cores(ssh)
        
        # 获取内存大小（MB）
        memory_size = get_memory_size(ssh)
        
        # 获取磁盘大小（GB）
        disk_size = get_disk_size(ssh)
        
        ssh.close()
        
        return {
            'success': True,
            'data': {
                'os_type': os_type,
                'cpu_cores': cpu_cores,
                'memory_size': memory_size,
                'disk_size': disk_size
            }
        }
    except Exception as e:
        if ssh:
            ssh.close()
        return {
            'success': False,
            'error': str(e)
        }


def execute_command(ssh, command):
    """执行SSH命令并返回输出"""
    try:
        stdin, stdout, stderr = ssh.exec_command(command)
        output = stdout.read().decode('utf-8').strip()
        error = stderr.read().decode('utf-8').strip()
        if error:
            return None
        return output
    except:
        return None


def get_os_info(ssh):
    """获取操作系统信息"""
    # 尝试从 /etc/os-release 获取
    os_release = execute_command(ssh, "cat /etc/os-release 2>/dev/null | grep PRETTY_NAME | cut -d'=' -f2 | tr -d '\"'")
    if os_release:
        return os_release
    
    # 尝试从 lsb_release 获取
    lsb_release = execute_command(ssh, "lsb_release -d 2>/dev/null | cut -f2")
    if lsb_release:
        return lsb_release
    
    # 尝试从 /etc/redhat-release 获取 (CentOS/RedHat)
    redhat_release = execute_command(ssh, "cat /etc/redhat-release 2>/dev/null")
    if redhat_release:
        return redhat_release
    
    # 尝试从 uname 获取基本信息
    uname = execute_command(ssh, "uname -s")
    if uname:
        return uname
    
    return "Unknown"


def get_cpu_cores(ssh):
    """获取CPU核心数"""
    # 使用 nproc 命令
    nproc = execute_command(ssh, "nproc 2>/dev/null")
    if nproc and nproc.isdigit():
        return int(nproc)
    
    # 从 /proc/cpuinfo 获取
    cpu_count = execute_command(ssh, "cat /proc/cpuinfo 2>/dev/null | grep -c processor")
    if cpu_count and cpu_count.isdigit():
        return int(cpu_count)
    
    # 从 sysctl 获取 (macOS)
    sysctl = execute_command(ssh, "sysctl -n hw.ncpu 2>/dev/null")
    if sysctl and sysctl.isdigit():
        return int(sysctl)
    
    return None


def get_memory_size(ssh):
    """获取内存大小（MB）"""
    # Linux: 从 /proc/meminfo 获取
    meminfo = execute_command(ssh, "cat /proc/meminfo 2>/dev/null | grep MemTotal | awk '{print $2}'")
    if meminfo and meminfo.isdigit():
        # MemTotal 单位是 KB，转换为 MB
        return int(int(meminfo) / 1024)
    
    # 使用 free 命令
    free_m = execute_command(ssh, "free -m 2>/dev/null | grep Mem | awk '{print $2}'")
    if free_m and free_m.isdigit():
        return int(free_m)
    
    # macOS: 使用 sysctl
    sysctl_mem = execute_command(ssh, "sysctl -n hw.memsize 2>/dev/null")
    if sysctl_mem and sysctl_mem.isdigit():
        # hw.memsize 单位是字节，转换为 MB
        return int(int(sysctl_mem) / 1024 / 1024)
    
    return None


def get_disk_size(ssh):
    """获取磁盘大小（GB）- 根分区"""
    # 使用 df 命令获取根分区大小
    df_output = execute_command(ssh, "df -BG / 2>/dev/null | tail -1 | awk '{print $2}' | sed 's/G//'")
    if df_output:
        try:
            return int(float(df_output))
        except:
            pass
    
    # 另一种方式
    df_output2 = execute_command(ssh, "df -h / 2>/dev/null | tail -1 | awk '{print $2}'")
    if df_output2:
        try:
            # 处理单位 (G, T, M)
            size_str = df_output2.upper()
            if 'T' in size_str:
                return int(float(size_str.replace('T', '')) * 1024)
            elif 'G' in size_str:
                return int(float(size_str.replace('G', '')))
            elif 'M' in size_str:
                return int(float(size_str.replace('M', '')) / 1024)
        except:
            pass
    
    return None
